#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
获取xiaozhi-server的MQTT配置信息脚本
用于在不通过管理界面的情况下获取客户端ID、用户名、密码等信息
"""

import json
import base64
import hashlib
import hmac
import yaml
import os
import sys
from pathlib import Path

def load_config():
    """加载配置文件"""
    # 首先尝试加载 data/.config.yaml
    config_path = Path("main/xiaozhi-server/data/.config.yaml")
    if config_path.exists():
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f) or {}
    else:
        config = {}
    
    # 如果 data/.config.yaml 中没有某些配置，则从 config.yaml 中补充
    main_config_path = Path("main/xiaozhi-server/config.yaml")
    if main_config_path.exists():
        with open(main_config_path, 'r', encoding='utf-8') as f:
            main_config = yaml.safe_load(f) or {}
        
        # 合并配置，data/.config.yaml 优先
        def merge_config(base, override):
            for key, value in base.items():
                if key not in override:
                    override[key] = value
                elif isinstance(value, dict) and isinstance(override[key], dict):
                    merge_config(value, override[key])
        
        merge_config(main_config, config)
    
    return config

def generate_password_signature(content: str, secret_key: str) -> str:
    """生成MQTT密码签名"""
    try:
        hmac_obj = hmac.new(
            secret_key.encode("utf-8"), content.encode("utf-8"), hashlib.sha256
        )
        signature = hmac_obj.digest()
        return base64.b64encode(signature).decode("utf-8")
    except Exception as e:
        print(f"生成MQTT密码签名失败: {e}")
        return ""

def generate_mqtt_config(device_id: str, device_model: str = "default"):
    """生成MQTT配置信息"""
    config = load_config()
    server_config = config.get("server", {})
    
    # 检查是否配置了MQTT网关
    mqtt_gateway_endpoint = server_config.get("mqtt_gateway")
    if not mqtt_gateway_endpoint or mqtt_gateway_endpoint == "null":
        print("❌ 未配置MQTT网关地址 (server.mqtt_gateway)")
        return None
    
    # 检查MQTT签名密钥
    signature_key = server_config.get("mqtt_signature_key")
    if not signature_key or signature_key == "null":
        print("❌ 未配置MQTT签名密钥 (server.mqtt_signature_key)")
        return None
    
    # 生成分组ID
    group_id = f"GID_{device_model}".replace(":", "_").replace(" ", "_")
    
    # 生成客户端ID
    mac_address_safe = device_id.replace(":", "_")
    mqtt_client_id = f"{group_id}@@@{mac_address_safe}@@@{mac_address_safe}"
    
    # 生成用户名（Base64编码的JSON）
    user_data = {"ip": "unknown"}
    user_data_json = json.dumps(user_data)
    username = base64.b64encode(user_data_json.encode("utf-8")).decode("utf-8")
    
    # 生成密码
    password = generate_password_signature(mqtt_client_id + "|" + username, signature_key)
    
    # 构建MQTT配置
    mqtt_config = {
        "endpoint": mqtt_gateway_endpoint,
        "client_id": mqtt_client_id,
        "username": username,
        "password": password,
        "publish_topic": "device-server",
        "subscribe_topic": f"devices/p2p/{mac_address_safe}",
    }
    
    return mqtt_config

def main():
    """主函数"""
    print("🔍 xiaozhi-server MQTT配置信息获取工具")
    print("=" * 50)
    
    # 检查是否在正确的目录
    if not os.path.exists("main/xiaozhi-server"):
        print("❌ 请在xiaozhi-esp32-server项目根目录下运行此脚本")
        sys.exit(1)
    
    # 获取设备ID
    if len(sys.argv) > 1:
        device_id = sys.argv[1]
    else:
        device_id = input("请输入设备ID (MAC地址，如 11:22:33:44:55:66): ").strip()
    
    if not device_id:
        print("❌ 设备ID不能为空")
        sys.exit(1)
    
    # 获取设备型号（可选）
    if len(sys.argv) > 2:
        device_model = sys.argv[2]
    else:
        device_model = input("请输入设备型号 (可选，默认为 default): ").strip() or "default"
    
    print(f"\n📱 设备ID: {device_id}")
    print(f"📱 设备型号: {device_model}")
    print("-" * 50)
    
    # 生成MQTT配置
    mqtt_config = generate_mqtt_config(device_id, device_model)
    
    if mqtt_config:
        print("✅ MQTT配置信息:")
        print(f"📡 MQTT服务器地址: {mqtt_config['endpoint']}")
        print(f"🆔 客户端ID: {mqtt_config['client_id']}")
        print(f"👤 用户名: {mqtt_config['username']}")
        print(f"🔐 密码: {mqtt_config['password']}")
        print(f"📤 发布主题: {mqtt_config['publish_topic']}")
        print(f"📥 订阅主题: {mqtt_config['subscribe_topic']}")
        
        print("\n" + "=" * 50)
        print("📋 完整JSON配置:")
        print(json.dumps(mqtt_config, indent=2, ensure_ascii=False))
        
        # 保存到文件
        output_file = f"mqtt_config_{device_id.replace(':', '_')}.json"
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(mqtt_config, f, indent=2, ensure_ascii=False)
        print(f"\n💾 配置已保存到: {output_file}")
    else:
        print("\n❌ 无法生成MQTT配置，请检查配置文件")
        print("\n💡 提示:")
        print("1. 确保在 data/.config.yaml 或 config.yaml 中配置了 server.mqtt_gateway")
        print("2. 确保在 data/.config.yaml 或 config.yaml 中配置了 server.mqtt_signature_key")

if __name__ == "__main__":
    main()